import os
import random
import hashlib
import numpy as np
from typing import List, Tuple, Dict, Any, Union
import hmac

# ---------------- 参数定义 ----------------
KYBER_N = 256  # 多项式维度
KYBER_Q = 3329  # 模数(质数)
KYBER_ROOT_OF_UNITY = 17  # 原根
KYBER_ROOT_OF_UNITY_INV = 3120  # 原根的逆元 (17^-1 mod 3329)

# 不同安全级别参数
KYBER_PARAMS = {
    "512": {"k": 2, "eta1": 3, "eta2": 2, "du": 10, "dv": 4},
    "768": {"k": 3, "eta1": 2, "eta2": 2, "du": 10, "dv": 4},
    "1024": {"k": 4, "eta1": 2, "eta2": 2, "du": 11, "dv": 5}
}

# ---------------- 基础运算 ----------------
def mod_q(x: int) -> int:
    """模q运算"""
    return x % KYBER_Q

def montgomery_reduce(a: int) -> int:
    """Montgomery约简"""
    return a % KYBER_Q

def barrett_reduce(a: int) -> int:
    """Barrett约简"""
    return a % KYBER_Q

def csubq(a: int) -> int:
    """条件约简"""
    return a - KYBER_Q if a >= KYBER_Q else a

# ---------------- 多项式操作 ----------------
class Poly:
    """多项式类"""
    def __init__(self, coeffs: List[int] = None):
        self.coeffs = [0] * KYBER_N if coeffs is None else coeffs[:KYBER_N]
        self._normalize()

    def _normalize(self):
        """标准化系数到[0, q-1]"""
        self.coeffs = [c % KYBER_Q for c in self.coeffs]

    def __add__(self, other):
        result = [(a + b) % KYBER_Q for a, b in zip(self.coeffs, other.coeffs)]
        return Poly(result)

    def __sub__(self, other):
        result = [(a - b) % KYBER_Q for a, b in zip(self.coeffs, other.coeffs)]
        return Poly(result)

    def __mul__(self, other):
        """标量乘法或多项式乘法"""
        if isinstance(other, int):
            return Poly([(c * other) % KYBER_Q for c in self.coeffs])
        
        # 多项式乘法
        result = [0] * KYBER_N
        for i in range(KYBER_N):
            for j in range(KYBER_N):
                idx = (i + j) % KYBER_N
                sign = -1 if (i + j) >= KYBER_N else 1
                result[idx] = (result[idx] + sign * self.coeffs[i] * other.coeffs[j]) % KYBER_Q
        return Poly(result)

    def compress(self, d: int) -> List[int]:
        """压缩多项式系数"""
        return [((c << d) + KYBER_Q//2) // KYBER_Q & ((1 << d) - 1) for c in self.coeffs]

    
    def decompress(compressed: List[int], d: int) -> 'Poly':
        """从压缩系数恢复多项式"""
        return Poly([(c * KYBER_Q + (1 << (d-1))) >> d for c in compressed])
    @staticmethod

    def to_bytes(self) -> bytes:
        """将多项式转换为字节表示"""
        # 每个系数用16位表示
        return bytes([(self.coeffs[i] >> (8 * j)) & 0xff for i in range(KYBER_N) for j in range(2)])

    @staticmethod
    def from_bytes(data: bytes) -> 'Poly':
        """从字节表示恢复多项式"""
        coeffs = []
        for i in range(KYBER_N):
            c = (data[2*i] << 8) | data[2*i + 1]
            coeffs.append(c % KYBER_Q)
        return Poly(coeffs)

# ---------------- -NTT变换 ----------------
def bit_reverse(x: int, bits: int) -> int:
    """位反转函数"""
    y = 0
    for i in range(bits):
        y = (y << 1) | (x & 1)
        x >>= 1
    return y

def ntt(p: Poly) -> Poly:
    """前向NTT变换"""
    n = KYBER_N
    logn = 8  # log2(256)
    coeffs = p.coeffs.copy()
    
    # 位反转置换
    coeffs_rev = [0] * n
    for i in range(n):
        j = bit_reverse(i, logn)
        coeffs_rev[j] = coeffs[i]
    
    # NTT计算
    root = KYBER_ROOT_OF_UNITY
    roots = [pow(root, i, KYBER_Q) for i in range(n//2)]
    
    k = 1
    while k < n:
        for j in range(0, n, 2*k):
            for i in range(j, j + k):
                t = roots[(i - j) * (n // (2*k))] * coeffs_rev[i + k]
                coeffs_rev[i + k] = (coeffs_rev[i] - t) % KYBER_Q
                coeffs_rev[i] = (coeffs_rev[i] + t) % KYBER_Q
        k <<= 1
    
    return Poly(coeffs_rev)

def invntt(p: Poly) -> Poly:
    """逆NTT变换"""
    n = KYBER_N
    logn = 8  # log2(256)
    coeffs = p.coeffs.copy()
    
    # 位反转置换
    coeffs_rev = [0] * n
    for i in range(n):
        j = bit_reverse(i, logn)
        coeffs_rev[j] = coeffs[i]
    
    # 逆NTT计算
    root = KYBER_ROOT_OF_UNITY_INV
    roots = [pow(root, i, KYBER_Q) for i in range(n//2)]
    n_inv = pow(n, KYBER_Q - 2, KYBER_Q)  # n的逆元
    
    k = 1
    while k < n:
        for j in range(0, n, 2*k):
            for i in range(j, j + k):
                t = roots[(i - j) * (n // (2*k))] * coeffs_rev[i + k]
                coeffs_rev[i + k] = (coeffs_rev[i] - t) % KYBER_Q
                coeffs_rev[i] = (coeffs_rev[i] + t) % KYBER_Q
        k <<= 1
    
    # 乘以n的逆元
    for i in range(n):
        coeffs_rev[i] = (coeffs_rev[i] * n_inv) % KYBER_Q
    
    return Poly(coeffs_rev)

def poly_basemul(a: Poly, b: Poly) -> Poly:
    """NTT域中的点乘"""
    return Poly([(a.coeffs[i] * b.coeffs[i]) % KYBER_Q for i in range(KYBER_N)])

# ---------------- 随机采样 ----------------
def random_bytes(n: int) -> bytes:
    """生成随机字节"""
    return os.urandom(n)

def prf(seed: bytes, nonce: int, length: int) -> bytes:
    """伪随机函数"""
    h = hashlib.shake_128()
    h.update(seed)
    h.update(nonce.to_bytes(1, byteorder='little'))
    return h.digest(length)

def cbd(eta: int, seed: bytes, nonce: int) -> Poly:
    """中心二项分布采样"""
    # 使用确定性随机数生成器
    prf_output = prf(seed, nonce, KYBER_N * eta // 4)
    coeffs = []
    
    for i in range(KYBER_N):
        a = 0
        for j in range(eta):
            byte_idx = (i * eta + j) // 8
            bit_idx = (i * eta + j) % 8
            if byte_idx < len(prf_output):
                a += (prf_output[byte_idx] >> bit_idx) & 1
        b = eta - a
        coeffs.append((a - b) % KYBER_Q)
    
    return Poly(coeffs)

def uniform_poly(seed: bytes, nonce: int) -> Poly:
    """均匀分布采样多项式"""
    prf_output = prf(seed, nonce, KYBER_N * 2)  # 每个系数需要16位
    coeffs = []
    
    for i in range(KYBER_N):
        c = (prf_output[2*i] << 8) | prf_output[2*i + 1]
        coeffs.append(c % KYBER_Q)
    
    return Poly(coeffs)

def generate_matrix(seed: bytes, transpose: bool = False) -> List[List[Poly]]:
    """生成随机多项式矩阵A"""
    k = KYBER_PARAMS["768"]["k"]  # 默认使用Kyber768
    result = [[None for _ in range(k)] for _ in range(k)]

    for i in range(k):
        for j in range(k):
            # 转置情况下交换下标
            idx1, idx2 = (j, i) if transpose else (i, j)
            nonce = (256 * idx1 + idx2) % 256
            result[i][j] = uniform_poly(seed, nonce)

    return result

# ---------------- Kyber PKE ----------------
class KyberPKE:
    def __init__(self, security: str = "768"):
        """初始化Kyber PKE"""
        self.params = KYBER_PARAMS[security]
        self.k = self.params["k"]

    def keygen(self) -> Tuple[Dict, Dict]:
        """密钥生成"""
        # 生成随机种子
        rho = random_bytes(32)
        sigma = random_bytes(32)

        # 生成矩阵A
        A = generate_matrix(rho)

        # 采样s和e
        s = []
        e = []
        for i in range(self.k):
            s.append(cbd(self.params["eta1"], sigma, i))
            e.append(cbd(self.params["eta1"], sigma, i + self.k))

        # 转换到NTT域
        s_hat = [ntt(poly) for poly in s]
        e_hat = [ntt(poly) for poly in e]

        # 计算t = As + e
        t = []
        for i in range(self.k):
            t_i = Poly([0] * KYBER_N)
            for j in range(self.k):
                t_i = t_i + poly_basemul(A[i][j], s_hat[j])
            t_i = t_i + e_hat[i]
            t.append(t_i)

        # 压缩公钥
        t_compressed = [poly.compress(self.params["du"]) for poly in t]

        return {"t": t_compressed, "rho": rho}, {"s": s_hat}

    def encrypt(self, pk: Dict, m: bytes, coins: bytes = None) -> Dict:
        """加密算法"""
        if coins is None:
            coins = random_bytes(32)

        rho = pk["rho"]
        t_compressed = pk["t"]
        t = [Poly.decompress(comp, self.params["du"]) for comp in t_compressed]

        # 生成矩阵A
        A = generate_matrix(rho, transpose=True)

        # 从消息m创建多项式m
        m_poly = self._bytes_to_poly(m)

        # 采样r和e1, e2
        r = []
        e1 = []
        for i in range(self.k):
            r.append(cbd(self.params["eta1"], coins, i))
            e1.append(cbd(self.params["eta2"], coins, i + self.k))
        e2 = cbd(self.params["eta2"], coins, 2 * self.k)

        # 转换到NTT域
        r_hat = [ntt(poly) for poly in r]

        # 计算u = A^T r + e1
        u = []
        for i in range(self.k):
            u_i = Poly([0] * KYBER_N)
            for j in range(self.k):
                u_i = u_i + poly_basemul(A[i][j], r_hat[j])
            # 逆NTT后加e1
            u_i = invntt(u_i) + e1[i]
            u.append(u_i)

        # 计算v = t^T r + e2 + m⌈q/2⌉
        v = Poly([0] * KYBER_N)
        for i in range(self.k):
            v = v + poly_basemul(t[i], r_hat[i])
        v = invntt(v) + e2 + m_poly * (KYBER_Q // 2)

        # 压缩密文
        u_compressed = [poly.compress(self.params["du"]) for poly in u]
        v_compressed = v.compress(self.params["dv"])

        return {"u": u_compressed, "v": v_compressed}

    def decrypt(self, sk: Dict, ct: Dict) -> bytes:
        """解密算法"""
        u_compressed = ct["u"]
        v_compressed = ct["v"]

        # 解压缩密文
        u = [Poly.decompress(comp, self.params["du"]) for comp in u_compressed]
        # 修复：v是单个多项式，不是列表
        v = Poly.decompress(v_compressed, self.params["dv"])

        # 计算m = v - s^T u
        s_hat = sk["s"]
        m_poly = v

        for i in range(self.k):
            u_hat = ntt(u[i])
            m_poly = m_poly - invntt(poly_basemul(s_hat[i], u_hat))

        # 解码消息
        return self._poly_to_bytes(m_poly)

    def _bytes_to_poly(self, data: bytes) -> Poly:
        """将字节转换为多项式(每个系数取低1位)"""
        coeffs = []
        for i in range(KYBER_N):
            if i < len(data):
                coeffs.append(data[i] & 1)
            else:
                coeffs.append(0)
        return Poly(coeffs)

    def _poly_to_bytes(self, poly: Poly) -> bytes:
        """将多项式转换为字节(每个系数的低1位)"""
        bytes_data = bytearray()
        for i in range(0, KYBER_N, 8):
            byte = 0
            for j in range(8):
                if i + j < KYBER_N:
                    byte |= (poly.coeffs[i+j] & 1) << j
            bytes_data.append(byte)
        return bytes(bytes_data)

# ---------------- Kyber KEM ----------------
class KyberKEM:
    def __init__(self, security: str = "768"):
        """初始化Kyber KEM"""
        self.pke = KyberPKE(security)

    def keygen(self) -> Tuple[Dict, Dict]:
        """KEM密钥生成"""
        pk, sk_pke = self.pke.keygen()

        # 额外的sk组件
        z = random_bytes(32)
        h_pk = hashlib.sha3_256(str(pk).encode()).digest()

        sk = {
            "sk_pke": sk_pke,
            "pk": pk,
            "h_pk": h_pk,
            "z": z
        }

        return pk, sk

    def encaps(self, pk: Dict) -> Tuple[bytes, Dict]:
        """密钥封装"""
        # 生成随机消息
        m = random_bytes(32)

        # 计算哈希
        h_pk = hashlib.sha3_256(str(pk).encode()).digest()

        # 使用sha256代替sha3-512
        g_hash = hashlib.sha256(m + h_pk).digest()

        # 分割hash输出为随机数和共享密钥
        coins = g_hash[:16]
        shared_key_input = g_hash[16:]

        # 使用确定性随机数加密
        ct = self.pke.encrypt(pk, m, coins)

        # 计算共享密钥
        ct_bytes = str(ct).encode()
        shared_key = hashlib.sha256(shared_key_input + ct_bytes).digest()

        return shared_key, ct

    def decaps(self, sk: Dict, ct: Dict) -> bytes:
        """密钥解封装"""
        # 解构私钥
        sk_pke = sk["sk_pke"]
        pk = sk["pk"]
        h_pk = sk["h_pk"]
        z = sk["z"]

        # 解密获取消息m'
        m_prime = self.pke.decrypt(sk_pke, ct)

        # 重新计算coins
        g_hash = hashlib.sha256(m_prime + h_pk).digest()
        coins = g_hash[:16]
        shared_key_input = g_hash[16:]

        # 重新加密验证
        ct_prime = self.pke.encrypt(pk, m_prime, coins)

        # 使用常数时间比较
        ct_bytes = str(ct).encode()
        ct_prime_bytes = str(ct_prime).encode()
        ct_matches = hmac.compare_digest(ct_bytes, ct_prime_bytes)

        # 如果密文匹配，使用正常密钥，否则使用伪随机密钥
        if ct_matches:
            return hashlib.sha256(shared_key_input + ct_bytes).digest()
        else:
            return hashlib.sha256(z + ct_bytes).digest()

# ---------------- 示例 ----------------
def demo_kyber():
    """Kyber演示"""
    print("=== CRYSTALS-Kyber 演示 ===")

    # 初始化Kyber-768
    kem = KyberKEM("768")

    # 密钥生成
    print("1.生成密钥")
    pk, sk = kem.keygen()

    # 密钥封装
    print("2.密钥封装")
    shared_secret_a, ct = kem.encaps(pk)

    # 密钥解封装
    print("3.密钥解封装")
    shared_secret_b = kem.decaps(sk, ct)

    # 验证
    print("共享密钥 A:", shared_secret_a.hex()[:20] + "...")
    print("共享密钥 B:", shared_secret_b.hex()[:20] + "...")
    print("密钥是否匹配:", shared_secret_a == shared_secret_b)

if __name__ == "__main__":
    demo_kyber()    